package locutor.draft.r

trait Rewriter {
  type Term = AnyRef
  
  abstract class Strategy extends (Term => Option[Term]) { p =>
    def apply (r : Term) : Option[Term]
    
    def <* (q : => Strategy) : Strategy = new Strategy {
      def apply (t1 : Term) = p(t1) match {
        case Some (t2) => q (t2)
        case None      => None
      }
    }
    
    def <+ (q : => Strategy) : Strategy = new Strategy {
      def apply (t1 : Term) = p(t1) match {
        case Some (t2) => Some (t2)
        case None      => q (t1)
      }
    }

    def + (q : => Strategy) : PlusStrategy = new PlusStrategy (p, q)

    def < (lr : => PlusStrategy) : Strategy = new Strategy {
      def apply (t1 : Term) = p(t1) match {
        case Some (t2) => lr.lhs (t2)
        case None      => lr.rhs (t1)
      }
    }
  }

  class PlusStrategy (p : => Strategy, q : => Strategy) extends Strategy {
    val lhs = p
    val rhs = q
    def apply (t : Term) = (p <+ q) (t)
  }

  def strategyf (f : Term => Option[Term]) : Strategy = new Strategy { def apply (t : Term) = f(t) }

  def strategy (f : PartialFunction[Term,Option[Term]]) : Strategy = new Strategy {
    def apply (t : Term) = if (f isDefinedAt t) f(t) else None
  }  

  def rulef (f : Term => Term) : Strategy = strategyf (t => Some(f(t)))

  def rule (f : PartialFunction[Term,Term]) : Strategy = new Strategy {
    def apply (t : Term) = if (f isDefinedAt t) Some(f(t)) else None
  }

  implicit def termToStrategy (t : Term) : Strategy = strategyf (_ => Some (t))

  def queryf[T] (f : Term => T) : Strategy = new Strategy {
    def apply (t : Term) = {
      val v = f (t)
      Some (t)
    }
  }

  def query[T] (f : PartialFunction[Term,T]) : Strategy = new Strategy {
    def apply (t : Term) = {
      if (f isDefinedAt t) f(t) else ()
      Some(t)
    }
  }

  val failure : Strategy = strategyf (_ => None)

  val id : Strategy = strategyf (t => Some (t))

  def term (t : Term) : Strategy =
    rule {
      case `t` => t
  }

  object Term {
    def unapply (t : Any) : Option[(Any,Seq[Any])] = t match {
      case p : Product =>
        val cs = for (i <- 0 until p.productArity) yield p.productElement (i)
        Some ((p, cs))
      case _ => Some ((t, Nil))
    }
  }

  def para[T] (f : (Any, Seq[T]) => T) : Any => T = { case Term (t, ts) => f (t, ts.map (para (f))) }

  /** General product duplication function.
   * 
   *  Returns a product that applies the same constructor as the product t, but with the given children
   *  instead of t's children.  Fails if a constructor cannot be found or
   *  if one of the children is not of the appropriate type. */
  private def dup (t : Product, children : Array[AnyRef]) : Product = {
    val ctor = (t.getClass.getConstructors())(0)
    try {
      ctor.newInstance (children : _*).asInstanceOf[Product]
    } catch {
      case e : java.lang.ClassCastException =>
        error ("dup cast failed: " + t)
      case e : IllegalArgumentException =>
        error ("dup illegal arguments: " + ctor + " " + children.deepMkString (",") + 
                 " (expects " + ctor.getParameterTypes.length + ")")
      }
    }

  /** Make an arbitrary value into a term child, checking that it worked properly.
   *  Object references will be returned unchanged; other values will be boxed. */
  private def makechild (child : Any) : AnyRef =
    try {
      return child.asInstanceOf[AnyRef]
    } catch {
      case e : ClassCastException => error ("makechild: can't cast child: " + child + " " + e)
    }

  def child (i : Int, s : Strategy) : Strategy = new Strategy {
    def apply (t : Term) : Option[Term] = t match {
      case p : Product =>
        val numchildren = p.productArity
        if ((i < 1) || (i > numchildren))
          None
        else p.productElement (i-1) match {
          case ct : Term =>
            val children = new Array[AnyRef](numchildren)
            for (j <- 0 until numchildren)
              children (j) = makechild (p.productElement (j))
            s (ct) match {
              case Some (ti) => children (i-1) = ti
              case None      => return None
            }
            val ret = dup (p, children)
            Some (ret)
          case ci => None
        }
      case a => None
    }
  }

  /** Traversal to all children.
   *  
   *  Construct a strategy that applies s to all term children of the subject term in left-to-right order.
   *  If s succeeds on all of the children, then succeed, forming a new term from the constructor
   *  of the original term and the result of s for each child.  If s fails on any child, fail. */
  def all (s : => Strategy) : Strategy = new Strategy {
    def apply (t : Term) : Option[Term] = t match {
      case p : Product =>
        val numchildren = p.productArity
        if (numchildren == 0)
          Some (t)
        else {
          val children = new Array[AnyRef](numchildren)
          for (i <- 0 until numchildren)
            p.productElement(i) match {
              case ct : Term =>
                s(ct) match {
                  case Some (ti) => children(i) = ti
                  case None      => return None
                }
              case ci =>
                // Child is not a term, don't try to transform it
                children(i) = makechild(ci)
            }          
          val ret = dup(p, children)
          Some(ret)
        }
      case a => Some(a)
    }
  }

  def one (s : => Strategy) : Strategy = new Strategy {
    def apply (t : Term) : Option[Term] = t match {
      case p : Product =>
        val numchildren = p.productArity
        for (i <- 0 until numchildren) {
          p.productElement (i) match {
            case ct : Term =>
              s (ct) match {
                case Some (nct) => {
                  val children = new Array[AnyRef] (numchildren)
                  for (j <- 0 until i)
                    children (j) = makechild (p.productElement (j))
                  children (i) = nct
                  for (j <- i + 1 until numchildren)
                    children (j) = makechild (p.productElement (j))
                  val ret = dup (p, children)
                  return Some (ret)
                }
                case None =>
                  // Transformation failed, this child stays unchanged
              }
            case _ =>
              // Child is not a term, don't try to transform it
          }
        }
        None
      case _ => None
    }
  }

  /** Traversal to as many children as possible, but at least one.
   * 
   *  Construct a strategy that applies s to the term children of the subject term in
   *  left-to-right order.  If s succeeds on any of the children, then succeed,
   *  forming a new term from the constructor of the original term and the result
   *  of s for each succeeding child, with other children unchanged.  In the event
   *  that the strategy fails on all children, then fail.*/
  def some (s : => Strategy) : Strategy = new Strategy {
    def apply (t : Term) : Option[Term] = t match {
      case p : Product =>
        val numchildren = p.productArity
        if (numchildren == 0)
          None
        else {
          val children = new Array[AnyRef](numchildren)
          var success = false
          for (i <- 0 until numchildren)
            p.productElement (i) match {
              case ct : Term =>
                s (ct) match {
                  case Some (ti) =>
                    children (i) = ti
                    success = true
                  case None      =>
                    // Child failed, don't transform it
                    children (i) = ct
                }
              case ci =>
                // Child is not a term, don't try to transform it
                children (i) = makechild (ci)
            }          
          if (success) {
            val ret = dup (p, children)
            Some (ret)
          } else None
        }
      case _ => None
    }
  }

  def rewrite[T <: Term] (s : => Strategy) (t : T) : T = s(t) match {
    case Some(t1) => t1.asInstanceOf[T]
    case None     => t
  }

  def collects[T] (f : PartialFunction[Term,T]) : Term => Set[T] = (t : Term) => {
    var collection = Set[T]()
    def collect = (v : T) => collection += v
    (everywheretd (query (f andThen collect))) (t)
    collection
  }

  def collectl[T] (f : PartialFunction[Term,T]) : Term => List[T] = (t : Term) => {
    var collection = List[T]()
    def collect = (v : T) => collection = collection ::: List (v)
    (everywheretd (query (f andThen collect))) (t)
    collection
  }

  def count (f : PartialFunction[Term,Int]) : Term => Int = (t : Term) => {
    var total = 0
    def count = (v : Int) => total += v
    (everywheretd (query (f andThen count))) (t)
    total
  }

  def attempt (s : => Strategy) : Strategy = s <+ id

  def repeat (s : => Strategy) : Strategy = attempt (s <* repeat (s))

  def repeat (s : => Strategy, c : => Strategy) : Strategy = (s <* repeat (s, c)) <+ c

  def repeat (s : => Strategy, n : Int) : Strategy = if (n == 0) id else s <* repeat (s, n - 1)

  def repeat1 (s : => Strategy, c : => Strategy) : Strategy = s <* (repeat1 (s, c) <+ c)

  def repeat1 (s : => Strategy) : Strategy = repeat1 (s, id)

  def repeatuntil (s : => Strategy, c : => Strategy) : Strategy = s <* (c <+ repeatuntil (s, c))

  def loop (c : => Strategy, s : => Strategy) : Strategy = attempt (c <* s <* loop (c, s))

  def loopnot (c : => Strategy, s : => Strategy) : Strategy = c <+ (s <* loopnot (c, s))

  def doloop (s : => Strategy, c : => Strategy) : Strategy = s <* loop (c, s)

  def loopiter (i : => Strategy, c : => Strategy, s : => Strategy) : Strategy = i <* loopnot (c, s)

  def loopiter (s : Int => Strategy, low : Int, up : Int) : Strategy =
    if (low <= up)
      s (low) <* loopiter (s, low + 1, up)
    else id

  def not (s : => Strategy) : Strategy = s < failure + id

  def where (s : => Strategy) : Strategy = strategyf (t => (s <* t) (t))

  def test (s : => Strategy) : Strategy = where (s)

  def breadthfirst (s : => Strategy) : Strategy = all (s) <* all (breadthfirst (s))

  def topdown (s : => Strategy) : Strategy = s <* all (topdown (s))

  def topdownS (s : => Strategy, stop : Strategy => Strategy) : Strategy =
    s <* (stop (topdownS (s, stop)) <+ all (topdownS (s, stop)))

  /** Construct a strategy that applies s in a bottom-up, postfix fashion to the subject term. */
  def bottomup (s : => Strategy) : Strategy = all (bottomup (s)) <* s

  def bottomupS (s : => Strategy, stop : Strategy => Strategy) : Strategy =
    (stop (bottomupS (s, stop)) <+ all (bottomupS (s, stop))) <* s

  def downup (s : => Strategy) : Strategy = s <* all (downup (s)) <* s

  def downup (s1 : => Strategy, s2 : => Strategy) : Strategy = s1 <* all (downup (s1, s2)) <* s2

  def downupS (s : => Strategy, stop : Strategy => Strategy) : Strategy =
    s <* (stop (downupS (s, stop)) <+ all (downupS (s, stop))) <* s

  def downupS (s1 : => Strategy, s2 : => Strategy, stop : Strategy => Strategy) : Strategy =
    s1 <* (stop (downupS (s1, s2, stop)) <+ all (downupS (s1, s2, stop))) <* s2

  def dontstop (s : => Strategy) : Strategy = failure

  def oncetd (s : => Strategy) : Strategy = s <+ one (oncetd (s))

  def oncebu (s : => Strategy) : Strategy = one (oncebu (s)) <+ s

  def sometd (s : => Strategy) : Strategy = s <+ some (sometd (s))

  def somebu (s : => Strategy) : Strategy = some (somebu (s)) <+ s

  def outermost (s : => Strategy) : Strategy = repeat (oncetd (s))

  def innermost (s : => Strategy) : Strategy = bottomup (attempt (s <* innermost (s)))

  def innermost2 (s : => Strategy) : Strategy = repeat (oncebu (s))

  /** Construct a strategy that applies s repeatedly to subterms until it fails on all of them.
   */
  def reduce (s : => Strategy) : Strategy = {
    def x : Strategy = some(x) + s
    repeat (x)
  }

  def alltd (s : => Strategy) : Strategy = s <+ all (alltd (s))

  def alldownup2 (s1 : => Strategy, s2 : => Strategy) : Strategy = (s1 <+ all (alldownup2 (s1, s2))) <* s2

  def alltdfold (s1 : => Strategy, s2 : => Strategy) : Strategy =
    s1 <+ (all (alltdfold (s1, s2)) <* s2)

  def somedownup (s : => Strategy) : Strategy =
    (s <* all (somedownup (s)) <* (attempt (s))) <+ (some (somedownup (s)) <+ attempt (s))

  def manybu (s : Strategy) : Strategy = some (manybu (s)) <* (attempt (s) <+ s)

  def manytd (s : Strategy) : Strategy = s <* all (attempt (manytd (s))) <+ some (manytd (s))

  val eq : Strategy =
    rule {
      case t @ (x, y) if x == y => t
    }

  val equal : Strategy = eq

  val issubterm : Strategy =
    strategy {
      case (x : Term, y : Term) => where (oncetd (term (x))) (y)
    }

  val ispropersubterm : Strategy = not (eq) <* issubterm

  val issuperterm : Strategy =
    strategy {
      case (x, y) => issubterm (y, x)
    }

  val ispropersuperterm : Strategy = not (eq) <* issuperterm

  val isleaf : Strategy = all (failure)

  def leaves (s : => Strategy, isleaf : => Strategy) : Strategy = (isleaf <* s) <+ all (leaves (s, isleaf))

  def leaves (s : => Strategy, isleaf : => Strategy, skip : Strategy => Strategy) : Strategy =
    (isleaf <* s) <+ skip (leaves (s, isleaf)) <+ all (leaves (s, isleaf))

  val isinnernode : Strategy = one (id)

  /** Construct a strategy that applies s at every term in a bottom-up fashion
   *  regardless of failure.  (Sub-)terms for which the strategy fails are left unchanged. */
  def everywherebu (s : => Strategy) : Strategy = bottomup (attempt (s))

  def everywheretd (s : => Strategy) : Strategy = topdown (attempt (s))

  def restore (s : => Strategy, rest : => Strategy) : Strategy = s <+ (rest <* failure)

  def restorealways (s : => Strategy, rest : => Strategy) : Strategy = s < rest + (rest <* failure)

  def lastly (s : => Strategy, f : => Strategy) : Strategy = s < where (f) + (where (f) <* failure)

  def ior (s1 : => Strategy, s2 : => Strategy) : Strategy = (s1 <* attempt (s2)) <+ s2

  def or (s1 : => Strategy, s2 : => Strategy) : Strategy = where (s1) < attempt (test (s2)) + test (s2)

  def and (s1 : => Strategy, s2 : => Strategy) : Strategy = where (s1) < test (s2) + (test (s2) <* failure)

}
